/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/* This file copied from Hadoop's security branch,
 * with the following changes:
 * 1. package changed from org.apache.hadoop.security to
 *    org.apache.zookeeper.server.auth.
 * 2. Usage of Hadoop's Configuration class removed since
 *    it is not available in Zookeeper: instead, system property
 *    "zookeeper.security.auth_to_local" is used.
 */

package org.apache.zookeeper.server.auth;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import sun.security.krb5.Config;
import sun.security.krb5.KrbException;

/**
 * This class implements parsing and handling of Kerberos principal names. In
 * particular, it splits them apart and translates them down into local
 * operating system names.
 */
public class KerberosName {
	@SuppressWarnings("serial")
	public static class BadFormatString extends IOException {
		BadFormatString(String msg) {
			super(msg);
		}

		BadFormatString(String msg, Throwable err) {
			super(msg, err);
		}
	}

	@SuppressWarnings("serial")
	public static class NoMatchingRule extends IOException {
		NoMatchingRule(String msg) {
			super(msg);
		}
	}

	/**
	 * An encoding of a rule for translating kerberos names.
	 */
	private static class Rule {
		/**
		 * Replace the numbered parameters of the form $n where n is from 1 to the
		 * length of params. Normal text is copied directly and $n is replaced by the
		 * corresponding parameter.
		 * 
		 * @param format the string to replace parameters again
		 * @param params the list of parameters
		 * @return the generated string with the parameter references replaced.
		 * @throws BadFormatString
		 */
		static String replaceParameters(String format, String[] params) throws BadFormatString {
			Matcher match = parameterPattern.matcher(format);
			int start = 0;
			StringBuilder result = new StringBuilder();
			while (start < format.length() && match.find(start)) {
				result.append(match.group(1));
				String paramNum = match.group(3);
				if (paramNum != null) {
					try {
						int num = Integer.parseInt(paramNum);
						if (num < 0 || num > params.length) {
							throw new BadFormatString("index " + num + " from " + format
									+ " is outside of the valid range 0 to " + (params.length - 1));
						}
						result.append(params[num]);
					} catch (NumberFormatException nfe) {
						throw new BadFormatString("bad format in username mapping in " + paramNum, nfe);
					}

				}
				start = match.end();
			}
			return result.toString();
		}

		/**
		 * Replace the matches of the from pattern in the base string with the value of
		 * the to string.
		 * 
		 * @param base   the string to transform
		 * @param from   the pattern to look for in the base string
		 * @param to     the string to replace matches of the pattern with
		 * @param repeat whether the substitution should be repeated
		 * @return
		 */
		static String replaceSubstitution(String base, Pattern from, String to, boolean repeat) {
			Matcher match = from.matcher(base);
			if (repeat) {
				return match.replaceAll(to);
			} else {
				return match.replaceFirst(to);
			}
		}

		private final String format;
		private final Pattern fromPattern;
		private final boolean isDefault;
		private final Pattern match;
		private final int numOfComponents;

		private final boolean repeat;

		private final String toPattern;

		Rule() {
			isDefault = true;
			numOfComponents = 0;
			format = null;
			match = null;
			fromPattern = null;
			toPattern = null;
			repeat = false;
		}

		Rule(int numOfComponents, String format, String match, String fromPattern, String toPattern, boolean repeat) {
			isDefault = false;
			this.numOfComponents = numOfComponents;
			this.format = format;
			this.match = match == null ? null : Pattern.compile(match);
			this.fromPattern = fromPattern == null ? null : Pattern.compile(fromPattern);
			this.toPattern = toPattern;
			this.repeat = repeat;
		}

		/**
		 * Try to apply this rule to the given name represented as a parameter array.
		 * 
		 * @param params first element is the realm, second and later elements are are
		 *               the components of the name "a/b@FOO" -> {"FOO", "a", "b"}
		 * @return the short name if this rule applies or null
		 * @throws IOException throws if something is wrong with the rules
		 */
		String apply(String[] params) throws IOException {
			String result = null;
			if (isDefault) {
				if (defaultRealm.equals(params[0])) {
					result = params[1];
				}
			} else if (params.length - 1 == numOfComponents) {
				String base = replaceParameters(format, params);
				if (match == null || match.matcher(base).matches()) {
					if (fromPattern == null) {
						result = base;
					} else {
						result = replaceSubstitution(base, fromPattern, toPattern, repeat);
					}
				}
			}
			if (result != null && nonSimplePattern.matcher(result).find()) {
				throw new NoMatchingRule("Non-simple name " + result + " after auth_to_local rule " + this);
			}
			return result;
		}

		@Override
		public String toString() {
			StringBuilder buf = new StringBuilder();
			if (isDefault) {
				buf.append("DEFAULT");
			} else {
				buf.append("RULE:[");
				buf.append(numOfComponents);
				buf.append(':');
				buf.append(format);
				buf.append(']');
				if (match != null) {
					buf.append('(');
					buf.append(match);
					buf.append(')');
				}
				if (fromPattern != null) {
					buf.append("s/");
					buf.append(fromPattern);
					buf.append('/');
					buf.append(toPattern);
					buf.append('/');
					if (repeat) {
						buf.append('g');
					}
				}
			}
			return buf.toString();
		}
	}

	private static String defaultRealm;

	private static Config kerbConf;

	/**
	 * A pattern that matches a Kerberos name with at most 2 components.
	 */
	private static final Pattern nameParser = Pattern.compile("([^/@]*)(/([^/@]*))?@([^/@]*)");

	/**
	 * A pattern that recognizes simple/non-simple names.
	 */
	private static final Pattern nonSimplePattern = Pattern.compile("[/@]");

	/**
	 * A pattern that matches a string with out '$' and then a single parameter with
	 * $n.
	 */
	private static Pattern parameterPattern = Pattern.compile("([^$]*)(\\$(\\d*))?");

	/**
	 * A pattern for parsing a auth_to_local rule.
	 */
	private static final Pattern ruleParser = Pattern
			.compile("\\s*((DEFAULT)|(RULE:\\[(\\d*):([^\\]]*)](\\(([^)]*)\\))?" + "(s/([^/]*)/([^/]*)/(g)?)?))");
	/**
	 * The list of translation rules.
	 */
	private static List<Rule> rules;

	static {
		try {
			kerbConf = Config.getInstance();
			defaultRealm = kerbConf.getDefaultRealm();
		} catch (KrbException ke) {
			if ((System.getProperty("zookeeper.requireKerberosConfig") != null)
					&& (System.getProperty("zookeeper.requireKerberosConfig").equals("true"))) {
				throw new IllegalArgumentException("Can't get Kerberos configuration", ke);
			}
		}
		try {
			// setConfiguration() will work even if the above try() fails due
			// to a missing Kerberos configuration (unless zookeeper.requireKerberosConfig
			// is set to true, which would not allow execution to reach here due to the
			// throwing of an IllegalArgumentException above).
			setConfiguration();
		} catch (IOException e) {
			throw new IllegalArgumentException("Could not configure Kerberos principal name mapping.");
		}
	}

	public static void main(String[] args) throws Exception {
		for (String arg : args) {
			KerberosName name = new KerberosName(arg);
			System.out.println("Name: " + name + " to " + name.getShortName());
		}
	}

	static List<Rule> parseRules(String rules) {
		List<Rule> result = new ArrayList<Rule>();
		String remaining = rules.trim();
		while (remaining.length() > 0) {
			Matcher matcher = ruleParser.matcher(remaining);
			if (!matcher.lookingAt()) {
				throw new IllegalArgumentException("Invalid rule: " + remaining);
			}
			if (matcher.group(2) != null) {
				result.add(new Rule());
			} else {
				result.add(new Rule(Integer.parseInt(matcher.group(4)), matcher.group(5), matcher.group(7),
						matcher.group(9), matcher.group(10), "g".equals(matcher.group(11))));
			}
			remaining = remaining.substring(matcher.end());
		}
		return result;
	}

	static void printRules() throws IOException {
		int i = 0;
		for (Rule r : rules) {
			System.out.println(++i + " " + r);
		}
	}

	/**
	 * Set the static configuration to get the rules.
	 * 
	 * @param conf the new configuration
	 * @throws IOException
	 */
	public static void setConfiguration() throws IOException {
		String ruleString = System.getProperty("zookeeper.security.auth_to_local", "DEFAULT");
		rules = parseRules(ruleString);
	}

	/** The second component of the name. It may be null. */
	private final String hostName;

	/** The realm of the name. */
	private final String realm;

	/** The first component of the name */
	private final String serviceName;

	/**
	 * Create a name from the full Kerberos principal name.
	 * 
	 * @param name
	 */
	public KerberosName(String name) {
		Matcher match = nameParser.matcher(name);
		if (!match.matches()) {
			if (name.contains("@")) {
				throw new IllegalArgumentException("Malformed Kerberos name: " + name);
			} else {
				serviceName = name;
				hostName = null;
				realm = null;
			}
		} else {
			serviceName = match.group(1);
			hostName = match.group(3);
			realm = match.group(4);
		}
	}

	/**
	 * Get the configured default realm.
	 * 
	 * @return the default realm from the krb5.conf
	 */
	public String getDefaultRealm() {
		return defaultRealm;
	}

	/**
	 * Get the second component of the name.
	 * 
	 * @return the second section of the Kerberos principal name, and may be null
	 */
	public String getHostName() {
		return hostName;
	}

	/**
	 * Get the realm of the name.
	 * 
	 * @return the realm of the name, may be null
	 */
	public String getRealm() {
		return realm;
	}

	/**
	 * Get the first component of the name.
	 * 
	 * @return the first section of the Kerberos principal name
	 */
	public String getServiceName() {
		return serviceName;
	}

	/**
	 * Get the translation of the principal name into an operating system user name.
	 * 
	 * @return the short name
	 * @throws IOException
	 */
	public String getShortName() throws IOException {
		String[] params;
		if (hostName == null) {
			// if it is already simple, just return it
			if (realm == null) {
				return serviceName;
			}
			params = new String[] { realm, serviceName };
		} else {
			params = new String[] { realm, serviceName, hostName };
		}
		for (Rule r : rules) {
			String result = r.apply(params);
			if (result != null) {
				return result;
			}
		}
		throw new NoMatchingRule("No rules applied to " + toString());
	}

	/**
	 * Put the name back together from the parts.
	 */
	@Override
	public String toString() {
		StringBuilder result = new StringBuilder();
		result.append(serviceName);
		if (hostName != null) {
			result.append('/');
			result.append(hostName);
		}
		if (realm != null) {
			result.append('@');
			result.append(realm);
		}
		return result.toString();
	}
}
